{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1_3oPpdBAduw"
      },
      "source": [
        "# Get started with semantic retrieval"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZFWzQEqNosrS"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/gemini-api/docs/semantic_retrieval\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on ai.google.dev</a>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemini-api/docs/semantic_retrieval.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/gemini-api/docs/semantic_retrieval.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AVW9cM0nuh1-"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Large Language Models (LLMs) can learn new abilities without directly being trained on them. However, LLMs have been known to \"hallucinate\" when tasked with providing responses for questions they have not been trained on. This is partly because LLMs are unaware of events after training. It is also very difficult to trace the sources from which LLMs draw their responses from. For reliable, scalable applications, it is important that an LLM provides responses that are grounded in facts and is able to cite its information sources.\n",
        "\n",
        "A common approach used to overcome these constraints is called Retrieval Augmented Generation (RAG), which augments the prompt sent to an LLM with relevant data retrieved from an external knowledge base through an Information Retrieval (IR) mechanism. The knowledge base can be your own corpora of documents, databases, or APIs.\n",
        "\n",
        "This notebook walks you through a workflow to improve an LLM's response by augmenting its knowledge with external text corpora and performing semantic information retrieval to answer questions using the Semantic Retriever and the Attributed Question & Answering (AQA) APIs of the Gemini API.\n",
        "\n",
        "Note: This API is currently in [beta](https://ai.google.dev/gemini-api/docs/api-versions) and is [only available in certain regions](https://ai.google.dev/gemini-api/docs/available-regions).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7jH3FO_BDua0"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uQwqEaFLHGlL"
      },
      "source": [
        "### Import the Gemini API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6029CKnWG75v"
      },
      "outputs": [],
      "source": [
        "# Install the Client library (Semantic Retriever is only supported for versions >0.4.0)\n",
        "!pip install -U google.ai.generativelanguage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G4D97GX5BRXb"
      },
      "source": [
        "## Authenticate\n",
        "\n",
        "The Semantic Retriever API lets you perform semantic search on your own data. Since it's **your data**, this needs stricter access controls than API keys. Authenticate with OAuth with [service accounts](#service-oauth) or through your [user credentials](#user-oauth).\n",
        "\n",
        "This quickstart uses a simplified authentication approach meant for a testing environment, and service account setups are typically easier to start from. For a production environment, learn about [authentication and authorization](https://developers.google.com/workspace/guides/auth-overview){:.external} before choosing the [access credentials](https://developers.google.com/workspace/guides/create-credentials#choose_the_access_credential_that_is_right_for_you){:.external} that are appropriate for your app.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eLjhFIOQ7_Dk"
      },
      "source": [
        "### Setup OAuth using service accounts {:#service-oauth}\n",
        "\n",
        "Follow the steps below to setup OAuth using service accounts:\n",
        "\n",
        "\n",
        "1. Enable the [Gemini API (Generative Language API)](https://console.cloud.google.com/flows/enableapi?apiid=generativelanguage.googleapis.com){:.external}.\n",
        "\n",
        "<img width=400 src=\"https://ai.google.dev/tutorials/images/semantic_retriever_enable_api.png\">\n",
        "\n",
        "2. Create the Service Account by following the [documentation](https://developers.google.com/identity/protocols/oauth2/service-account#creatinganaccount){:.external}.\n",
        "\n",
        " * After creating the service account, generate a service account key.\n",
        "\n",
        "<img width=400 src=\"https://ai.google.dev/tutorials/images/semantic_retriever_service_account.png\">\n",
        "\n",
        "3. Upload your service account file by using the file icon on the left sidebar, then the upload icon, as shown in the screenshot below.\n",
        "\n",
        " * Rename the uploaded file to `service_account_key.json` or change the variable `service_account_file_name` in the code below.\n",
        "\n",
        "<img width=400 src=\"https://ai.google.dev/tutorials/images/colab_upload.png\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "akwWUSrp8Bx2"
      },
      "outputs": [],
      "source": [
        "!pip install -U google-auth-oauthlib"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2jZmqVCj8FKa"
      },
      "outputs": [],
      "source": [
        "service_account_file_name = 'service_account_key.json'\n",
        "\n",
        "from google.oauth2 import service_account\n",
        "\n",
        "credentials = service_account.Credentials.from_service_account_file(service_account_file_name)\n",
        "\n",
        "scoped_credentials = credentials.with_scopes(\n",
        "    ['https://www.googleapis.com/auth/cloud-platform', 'https://www.googleapis.com/auth/generative-language.retriever'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qqjs2tvvicqq"
      },
      "source": [
        "Initialize the client library using the service account credentials."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P719DMtK8t-p"
      },
      "outputs": [],
      "source": [
        "import google.ai.generativelanguage as glm\n",
        "generative_service_client = glm.GenerativeServiceClient(credentials=scoped_credentials)\n",
        "retriever_service_client = glm.RetrieverServiceClient(credentials=scoped_credentials)\n",
        "permission_service_client = glm.PermissionServiceClient(credentials=scoped_credentials)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4EQJD2PWD56T"
      },
      "source": [
        "## Create a corpus {:#create-corpus}\n",
        "\n",
        "The Semantic Retriever API lets you define up to 5 custom text corpora per project. You can specify either of the following fields while defining your corpora:\n",
        "\n",
        "  * `name`: The `Corpus` resource name (ID). Must contain only a maximum of 40 alphanumeric characters. If the `name` is empty on creation, a unique name will be generated with a maximum length of 40 characters with a prefix from the `display_name` and a 12 character random suffix.\n",
        "  * `display_name`: The human-readable display name for the `Corpus`. Must contain only a maximum of 512 characters, including alphanumerics, spaces, and dashes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AaPZiXVwEDHZ"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "name: \"corpora/google-for-developers-blog-dqrtz8rs0jg\"\n",
            "display_name: \"Google for Developers Blog\"\n",
            "create_time {\n",
            "  seconds: 1713497533\n",
            "  nanos: 587977000\n",
            "}\n",
            "update_time {\n",
            "  seconds: 1713497533\n",
            "  nanos: 587977000\n",
            "}\n",
            "\n"
          ]
        }
      ],
      "source": [
        "example_corpus = glm.Corpus(display_name=\"Google for Developers Blog\")\n",
        "create_corpus_request = glm.CreateCorpusRequest(corpus=example_corpus)\n",
        "\n",
        "# Make the request\n",
        "create_corpus_response = retriever_service_client.create_corpus(create_corpus_request)\n",
        "\n",
        "# Set the `corpus_resource_name` for subsequent sections.\n",
        "corpus_resource_name = create_corpus_response.name\n",
        "print(create_corpus_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gQr7gDKVErdg"
      },
      "source": [
        "### Get the created corpus\n",
        "\n",
        "Use the `GetCorpusRequest` method to programmatically access the `Corpus` you created above. The value of the `name` parameter refers to the full resource name of the `Corpus` and is set in the cell above as `corpus_resource_name`. The expected format is `corpora/corpus-123`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8zhexURaEtVa"
      },
      "outputs": [],
      "source": [
        "get_corpus_request = glm.GetCorpusRequest(name=corpus_resource_name)\n",
        "\n",
        "# Make the request\n",
        "get_corpus_response = retriever_service_client.get_corpus(get_corpus_request)\n",
        "\n",
        "# Print the response\n",
        "print(get_corpus_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lbrLPNXRF2qP"
      },
      "source": [
        "## Create a document\n",
        "\n",
        "A `Corpus` can contain up to 10,000 `Document`s. You can specify either of the following fields while defining your documents:\n",
        "\n",
        "  * `name`: The `Document` resource name (ID). Must contain only a maximum of 40 characters (alphanumeric or dashes only). The ID cannot start or end with a\n",
        "  dash. If the name is empty on creation, a unique name will be derived from\n",
        "  `display_name` along with a 12 character random suffix.\n",
        "  * `display_name`: The human-readable display name. Must contain only a maximum of 512 characters, including alphanumerics, spaces, and dashes.\n",
        "\n",
        "`Document`s also support up to 20 user-specified `custom_metadata` fields, specified as key-value pairs. Custom metadata can be strings, lists of strings, or numeric. Note that lists of strings can support a maximum of 10 values and numeric values are represented as floating-point numbers in the API."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q7nwjPaGF_Nj"
      },
      "outputs": [],
      "source": [
        "# Create a document with a custom display name.\n",
        "example_document = glm.Document(display_name=\"Introducing Project IDX, An Experiment to Improve Full-stack, Multiplatform App Development\")\n",
        "\n",
        "# Add metadata.\n",
        "# Metadata also supports numeric values not specified here\n",
        "document_metadata = [\n",
        "    glm.CustomMetadata(key=\"url\", string_value=\"https://developers.googleblog.com/2023/08/introducing-project-idx-experiment-to-improve-full-stack-multiplatform-app-development.html\")]\n",
        "example_document.custom_metadata.extend(document_metadata)\n",
        "\n",
        "# Make the request\n",
        "# corpus_resource_name is a variable set in the \"Create a corpus\" section.\n",
        "create_document_request = glm.CreateDocumentRequest(parent=corpus_resource_name, document=example_document)\n",
        "create_document_response = retriever_service_client.create_document(create_document_request)\n",
        "\n",
        "# Set the `document_resource_name` for subsequent sections.\n",
        "document_resource_name = create_document_response.name\n",
        "print(create_document_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AiKINhAtGd8e"
      },
      "source": [
        "### Get the created document\n",
        "\n",
        "Use the `GetDocumentRequest` method to programmatically access the document you created above. The value of the `name` parameter refers to the full resource name of the document and is set in the cell above as `document_resource_name`. The expected format is `corpora/corpus-123/documents/document-123`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lTTc_gtWGfpe"
      },
      "outputs": [],
      "source": [
        "get_document_request = glm.GetDocumentRequest(name=document_resource_name)\n",
        "\n",
        "# Make the request\n",
        "# document_resource_name is a variable set in the \"Create a document\" section.\n",
        "get_document_response = retriever_service_client.get_document(get_document_request)\n",
        "\n",
        "# Print the response\n",
        "print(get_document_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OuuDtAYBOX55"
      },
      "source": [
        "## Ingest & Chunk a Document\n",
        "\n",
        "To improve the relevance of content returned by the vector database during semantic retrieval, break down large documents into smaller pieces or **chunks** while ingesting the document.\n",
        "\n",
        "A `Chunk` is a subpart of a `Document` that is treated as an independent unit for the purposes of vector representation and storage. A `Chunk` can have a maximum of 2043 tokens. A `Corpus` can have a maximum of 1 million `Chunk`s.\n",
        "\n",
        "Similar to `Document`s, `Chunks` also support up to 20 user-specified `custom_metadata` fields, specified as key-value pairs. Custom metadata can be strings, lists of strings, or numeric. Note that lists of strings can support a maximum of 10 values and numeric values are represented as floating-point numbers in the API.\n",
        "\n",
        "This guide uses Google's [Open Source HtmlChunker](https://github.com/google/labs-prototypes/tree/main/seeds/chunker-python){:.external}.\n",
        "\n",
        "Other chunkers you can use include [LangChain](https://python.langchain.com/docs/get_started/introduction){:.external} or [LlamaIndex](https://www.llamaindex.ai/){:.external}."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SkMzKWTbPBOU"
      },
      "source": [
        "### Ingest HTML and chunk via HtmlChunker"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nh7mwa7MPEZ9"
      },
      "outputs": [],
      "source": [
        "!pip install google-labs-html-chunker\n",
        "\n",
        "from google_labs_html_chunker.html_chunker import HtmlChunker\n",
        "\n",
        "from urllib.request import urlopen"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GUZYcNrEM-AK"
      },
      "source": [
        "Get the HTML DOM for a website. Here, the HTML is read directly, but it would\n",
        "be better to get HTML post-rendering to include Javascript-injected HTML\n",
        "such as `document.documentElement.innerHTML`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lgtl2Ow5PO94"
      },
      "outputs": [],
      "source": [
        "with(urlopen(\"https://developers.googleblog.com/2023/08/introducing-project-idx-experiment-to-improve-full-stack-multiplatform-app-development.html\")) as f:\n",
        "  html = f.read().decode(\"utf-8\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6oMLH2BINIde"
      },
      "source": [
        "Break down the text document into passages and create `Chunk`s from these passages. This step creates the `Chunk` objects themselves and the next section uploads them to the Semantic Retriever API."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z5G0oN4PNJEf"
      },
      "outputs": [],
      "source": [
        "# Chunk the file using HtmlChunker\n",
        "chunker = HtmlChunker(\n",
        "    max_words_per_aggregate_passage=200,\n",
        "    greedily_aggregate_sibling_nodes=True,\n",
        "    html_tags_to_exclude={\"noscript\", \"script\", \"style\"},\n",
        ")\n",
        "passages = chunker.chunk(html)\n",
        "print(passages)\n",
        "\n",
        "\n",
        "# Create `Chunk` entities.\n",
        "chunks = []\n",
        "for passage in passages:\n",
        "    chunk = glm.Chunk(data={'string_value': passage})\n",
        "    # Optionally, you can add metadata to a chunk\n",
        "    chunk.custom_metadata.append(glm.CustomMetadata(key=\"tags\",\n",
        "                                                    string_list_value=glm.StringList(\n",
        "                                                        values=[\"Google For Developers\", \"Project IDX\", \"Blog\", \"Announcement\"])))\n",
        "    chunk.custom_metadata.append(glm.CustomMetadata(key=\"chunking_strategy\",\n",
        "                                                    string_value=\"greedily_aggregate_sibling_nodes\"))\n",
        "    chunk.custom_metadata.append(glm.CustomMetadata(key = \"publish_date\",\n",
        "                                                    numeric_value = 20230808))\n",
        "    chunks.append(chunk)\n",
        "print(chunks)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIqgob2fUxeO"
      },
      "source": [
        "## Batch create chunks\n",
        "\n",
        "Create chunks in batches. You can specify a maximum of 100 chunks per batch request.\n",
        "\n",
        "Use `CreateChunk()` for single chunk creation.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iZX3oiOiU0YA"
      },
      "outputs": [],
      "source": [
        "# Option 1: Use HtmlChunker in the section above.\n",
        "# `chunks` is the variable set from the section above.\n",
        "create_chunk_requests = []\n",
        "for chunk in chunks:\n",
        "  create_chunk_requests.append(glm.CreateChunkRequest(parent=document_resource_name, chunk=chunk))\n",
        "\n",
        "# Make the request\n",
        "request = glm.BatchCreateChunksRequest(parent=document_resource_name, requests=create_chunk_requests)\n",
        "response = retriever_service_client.batch_create_chunks(request)\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xnn-IDSPGLcf"
      },
      "source": [
        "Alternatively, you can make chunks without using the HtmlChunker."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fE7_ueBPGMib"
      },
      "outputs": [],
      "source": [
        "# Add up to 100 CreateChunk requests per batch request.\n",
        "# document_resource_name is a variable set in the \"Create a document\" section.\n",
        "chunks = []\n",
        "chunk_1 = glm.Chunk(data={'string_value': \"Chunks support user specified metadata.\"})\n",
        "chunk_1.custom_metadata.append(glm.CustomMetadata(key=\"section\",\n",
        "                                                  string_value=\"Custom metadata filters\"))\n",
        "chunk_2 = glm.Chunk(data={'string_value': \"The maximum number of metadata supported is 20\"})\n",
        "chunk_2.custom_metadata.append(glm.CustomMetadata(key = \"num_keys\",\n",
        "                                                  numeric_value = 20))\n",
        "chunks = [chunk_1, chunk_2]\n",
        "create_chunk_requests = []\n",
        "for chunk in chunks:\n",
        "  create_chunk_requests.append(glm.CreateChunkRequest(parent=document_resource_name, chunk=chunk))\n",
        "\n",
        "# Make the request\n",
        "request = glm.BatchCreateChunksRequest(parent=document_resource_name, requests=create_chunk_requests)\n",
        "response = retriever_service_client.batch_create_chunks(request)\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r8kTaWguU1ws"
      },
      "source": [
        "### List `Chunk`s and get state\n",
        "\n",
        "Use the `ListChunksRequest` method to get all available `Chunk`s as a paginated list with a maximum size limit of 100 `Chunk`s per page, sorted in ascending order of `Chunk.create_time`. If you do not specify a limit, a maximum of 10 `Chunk`s are returned.\n",
        "\n",
        "Provide the `next_page_token` returned in the `ListChunksRequest` response as an argument to the next request to retrieve the next page. Note that when paginating, all other parameters provided to `ListChunks` must match the call that provided the page token.\n",
        "\n",
        "All `Chunk`s return a `state`. Use this to check the state of the `Chunks` before querying a `Corpus`. `Chunk` states include - `UNSPECIFIED`, `PENDING_PROCESSING`, `ACTIVE`, and `FAILED`. You can only query `ACTIVE` `Chunk`s."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DsNm8oJN6j18"
      },
      "outputs": [],
      "source": [
        "# Make the request\n",
        "request = glm.ListChunksRequest(parent=document_resource_name)\n",
        "list_chunks_response = retriever_service_client.list_chunks(request)\n",
        "for index, chunks in enumerate(list_chunks_response.chunks):\n",
        "  print(f'\\nChunk # {index + 1}')\n",
        "  print(f'Resource Name: {chunks.name}')\n",
        "  # Only ACTIVE chunks can be queried.\n",
        "  print(f'State: {glm.Chunk.State(chunks.state).name}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IRg1uAZSAaWS"
      },
      "source": [
        "## Ingest another document\n",
        "\n",
        "Add another `Document` via HtmlChunker and add filters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "quSpcAkRAX7L"
      },
      "outputs": [],
      "source": [
        "# Create a document with a custom display name.\n",
        "example_document = glm.Document(display_name=\"How it’s Made: Interacting with Gemini through multimodal prompting\")\n",
        "\n",
        "# Add document metadata.\n",
        "# Metadata also supports numeric values not specified here\n",
        "document_metadata = [\n",
        "    glm.CustomMetadata(key=\"url\", string_value=\"https://developers.googleblog.com/2023/12/how-its-made-gemini-multimodal-prompting.html\")]\n",
        "example_document.custom_metadata.extend(document_metadata)\n",
        "\n",
        "# Make the CreateDocument request\n",
        "# corpus_resource_name is a variable set in the \"Create a corpus\" section.\n",
        "create_document_request = glm.CreateDocumentRequest(parent=corpus_resource_name, document=example_document)\n",
        "create_document_response = retriever_service_client.create_document(create_document_request)\n",
        "\n",
        "# Set the `document_resource_name` for subsequent sections.\n",
        "document_resource_name = create_document_response.name\n",
        "print(create_document_response)\n",
        "\n",
        "# Chunks - add another webpage from Google for Developers\n",
        "with(urlopen(\"https://developers.googleblog.com/2023/12/how-its-made-gemini-multimodal-prompting.html\")) as f:\n",
        "  html = f.read().decode(\"utf-8\")\n",
        "\n",
        "# Chunk the file using HtmlChunker\n",
        "chunker = HtmlChunker(\n",
        "    max_words_per_aggregate_passage=100,\n",
        "    greedily_aggregate_sibling_nodes=False,\n",
        ")\n",
        "passages = chunker.chunk(html)\n",
        "\n",
        "# Create `Chunk` entities.\n",
        "chunks = []\n",
        "for passage in passages:\n",
        "    chunk = glm.Chunk(data={'string_value': passage})\n",
        "    chunk.custom_metadata.append(glm.CustomMetadata(key=\"tags\",\n",
        "                                                    string_list_value=glm.StringList(\n",
        "                                                        values=[\"Google For Developers\", \"Gemini API\", \"Blog\", \"Announcement\"])))\n",
        "    chunk.custom_metadata.append(glm.CustomMetadata(key=\"chunking_strategy\",\n",
        "                                                    string_value=\"no_aggregate_sibling_nodes\"))\n",
        "    chunk.custom_metadata.append(glm.CustomMetadata(key = \"publish_date\",\n",
        "                                                    numeric_value = 20231206))\n",
        "    chunks.append(chunk)\n",
        "\n",
        "# Make the request\n",
        "create_chunk_requests = []\n",
        "for chunk in chunks:\n",
        "  create_chunk_requests.append(glm.CreateChunkRequest(parent=document_resource_name, chunk=chunk))\n",
        "request = glm.BatchCreateChunksRequest(parent=document_resource_name, requests=create_chunk_requests)\n",
        "response = retriever_service_client.batch_create_chunks(request)\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "narzqqh0U0Ii"
      },
      "source": [
        "## Query the corpus\n",
        "\n",
        "Use the `QueryCorpusRequest` method to perform semantic search to get relevant passages.\n",
        "\n",
        "* `results_count`: Specify the number of passages to return. Maximum is 100. If unspecified, the API returns a maximum of 10 `Chunk`s.\n",
        "* `metadata_filters`: Filter by `chunk_metadata` or `document_metadata`. Each `MetadataFilter` needs to correspond to a unique key. Multiple `MetadataFilter` objects are joined by logical `AND`s. Similar metadata filter conditions are joined by logical `OR`s. Some examples:\n",
        "\n",
        "```\n",
        "(year >= 2020 OR year < 2010) AND (genre = drama OR genre = action)\n",
        "\n",
        "metadata_filter = [\n",
        "  {\n",
        "    key = \"document.custom_metadata.year\"\n",
        "    conditions = [\n",
        "      {int_value = 2020, operation = GREATER_EQUAL},\n",
        "      {int_value = 2010, operation = LESS}]\n",
        "  },\n",
        "  {\n",
        "    key = \"document.custom_metadata.genre\"\n",
        "    conditions = [\n",
        "      {string_value = \"drama\", operation = EQUAL},\n",
        "      {string_value = \"action\", operation = EQUAL}}]\n",
        "  }]\n",
        "```\n",
        "\n",
        "Note that only numeric values support \"AND\"s for the same key. String\n",
        "values only support \"OR\"s for the same key.\n",
        "\n",
        "```\n",
        "(\"Google for Developers\" in tags) and (20230314 > publish_date)\n",
        "\n",
        "metadata_filter = [\n",
        " {\n",
        "    key = \"chunk.custom_metadata.tags\"\n",
        "    conditions = [\n",
        "    {string_value = 'Google for Developers', operation = INCLUDES},\n",
        "  },\n",
        "  {\n",
        "    key = \"chunk.custom_metadata.publish_date\"\n",
        "    conditions = [\n",
        "    {numeric_value = 20230314, operation = GREATER_EQUAL}]\n",
        "  }]\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZtoGd0yz7wVg"
      },
      "outputs": [],
      "source": [
        "user_query = \"What is the purpose of Project IDX?\"\n",
        "results_count = 5\n",
        "\n",
        "# Add metadata filters for both chunk and document.\n",
        "chunk_metadata_filter = glm.MetadataFilter(key='chunk.custom_metadata.tags',\n",
        "                                           conditions=[glm.Condition(\n",
        "                                              string_value='Google For Developers',\n",
        "                                              operation=glm.Condition.Operator.INCLUDES)])\n",
        "\n",
        "# Make the request\n",
        "# corpus_resource_name is a variable set in the \"Create a corpus\" section.\n",
        "request = glm.QueryCorpusRequest(name=corpus_resource_name,\n",
        "                                 query=user_query,\n",
        "                                 results_count=results_count,\n",
        "                                 metadata_filters=[chunk_metadata_filter])\n",
        "query_corpus_response = retriever_service_client.query_corpus(request)\n",
        "print(query_corpus_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9yVhNmjkVER2"
      },
      "source": [
        "## Attributed Question-Answering\n",
        "\n",
        "Use the [`GenerateAnswer`](https://ai.google.dev/api/python/google/generativeai/protos/GenerateAnswerRequest) method to perform Attributed Question-Answering over your document, corpus, or a set of passages.\n",
        "\n",
        "Attributed Question-Answering (AQA) refers to answering questions grounded to a given context and providing attributions(s), while minimizing hallucination.\n",
        "\n",
        "`GenerateAnswer` provides several advantages over using an untuned LLM, in cases where AQA is desired:\n",
        " * The underlying model has been trained to return only answers that are grounded in the supplied context.\n",
        " * It identifies attributions (segments of the supplied context that contributed to the answer). Attributions enable the user to verify the answer.\n",
        " * It estimates the `answerable_probability` for a given (question, context) pair, which further empowers you to divert product behavior depending on how likely the returned answer is to be grounded and correct.\n",
        "\n",
        "Note: AQA currently only supports queries in English.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tPNpJdjJaYEz"
      },
      "source": [
        "### `answerable_probability` and the “I don’t know” problem\n",
        "\n",
        "In some instances, the best response to the question is in fact “I don’t know”. For example, if the provided context does not contain the answer to the question, then the question is considered to be “unanswerable”.\n",
        "\n",
        "The AQA model is highly adept at recognizing such cases. It can even distinguish between degrees of answerability and unanswerability.\n",
        "\n",
        "However, the `GenerateAnswer` API puts the final decision-making power in your hands by:\n",
        "* *Always* attempting to return a grounded answer - even when that answer is relatively unlikely to be grounded and correct.\n",
        "* Returning a value `answerable_probability` - The model's estimate of the probability that the answer is grounded and correct.\n",
        "\n",
        "A low `answerable_probability` may be explained by 1 or more of the following factors:\n",
        "\n",
        "* The model is not confident that its answer is correct.\n",
        "* The model is not confident that its answer is grounded in the cited passages; The answer may be derived instead from world knowledge. For example: `question=\"1+1=?\", passages=[\"2+2=4”]` → `answer=2, answerable_probability=0.02`\n",
        "* The model provided relevant information that did not completely answer the question. Example: `question=\"Is it available in my size?, passages=[\"Available in sizes 5-11\"]` → `answer=\"Yes it is available in sizes 5-11\", answerable_probability=0.03\"`\n",
        "* No well-formed question was asked in the GenerateAnswerRequest.\n",
        "\n",
        "Since a low `answerable_probability` indicates that the GenerateAnswerResponse.answer is likely wrong or ungrounded, **it is highly recommended to further process the response by inspecting `answerable_probability`**.\n",
        "\n",
        "When `answerable_probability` is low, some clients may wish to:\n",
        "* Display a message to the effect of \"couldn't answer that question\" to the end user.\n",
        "* Fall back to a general-purpose LLM that answers the question from world knowledge. The threshold and nature of such fallbacks will depend on individual use cases. A value of `answerable_probability` <= 0.5 is a good starting threshold.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aoUQ37Vqad0V"
      },
      "source": [
        "### AQA Helpful Tips\n",
        "\n",
        "For full API specifications, refer to the [`GenerateAnswerRequest` API Reference](https://ai.google.dev/api/python/google/generativeai/protos/GenerateAnswerRequest).\n",
        "\n",
        "* *Passage length*: Up to 300 tokens per passage are recommended.\n",
        "* *Passage sorting*:\n",
        "  * If you provide `GenerateAnswerRequest.inline_passages`, the passages should be sorted in decreasing order of relevance to the query. If the model's context length limit is exceeded, the last (least-relevant) passages will be omitted.\n",
        "  * If you provide `GenerateAnswerRequest.semantic_retriever`, then relevance sorting will be done automatically for you.\n",
        "* *Limitations*: The AQA model is specialized for question-answering. For other use cases such as creative writing, summarization, etc., please call a general-purpose model via GenerateContent.\n",
        "  * *Chat*: If the user input is known to be a question that may be answerable from a certain context, then AQA can answer chat queries. But if user input may be any type of entry, then a general-purpose model may be a better choice.\n",
        "* *Temperature*:\n",
        "  * Generally, a relatively low (~0.2) temperature is recommended for accurate AQA.\n",
        "  * If your use case relies on deterministic outputs, then set temperature=0.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SJr0vi7y_Uho"
      },
      "outputs": [],
      "source": [
        "user_query = \"What is the purpose of Project IDX?\"\n",
        "answer_style = \"ABSTRACTIVE\" # Or VERBOSE, EXTRACTIVE\n",
        "MODEL_NAME = \"models/aqa\"\n",
        "\n",
        "# Make the request\n",
        "# corpus_resource_name is a variable set in the \"Create a corpus\" section.\n",
        "content = glm.Content(parts=[glm.Part(text=user_query)])\n",
        "retriever_config = glm.SemanticRetrieverConfig(source=corpus_resource_name, query=content)\n",
        "req = glm.GenerateAnswerRequest(model=MODEL_NAME,\n",
        "                                contents=[content],\n",
        "                                semantic_retriever=retriever_config,\n",
        "                                answer_style=answer_style)\n",
        "aqa_response = generative_service_client.generate_answer(req)\n",
        "print(aqa_response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ack_zCIBHGkH"
      },
      "outputs": [],
      "source": [
        "# Get the metadata from the first attributed passages for the source\n",
        "chunk_resource_name = aqa_response.answer.grounding_attributions[0].source_id.semantic_retriever_chunk.chunk\n",
        "get_chunk_response = retriever_service_client.get_chunk(name=chunk_resource_name)\n",
        "print(get_chunk_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wjYUSa_5awjR"
      },
      "source": [
        "### More Options: AQA Using Inline Passages\n",
        "\n",
        "Alternatively, you can use the AQA endpoint directly, without using the Semantic Retriever API by passing `inline_passages`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AipBQbxqawPO"
      },
      "outputs": [],
      "source": [
        "user_query = \"What is AQA from Google?\"\n",
        "user_query_content = glm.Content(parts=[glm.Part(text=user_query)])\n",
        "answer_style = \"VERBOSE\" # or ABSTRACTIVE, EXTRACTIVE\n",
        "MODEL_NAME = \"models/aqa\"\n",
        "\n",
        "# Create the grounding inline passages\n",
        "grounding_passages = glm.GroundingPassages()\n",
        "passage_a = glm.Content(parts=[glm.Part(text=\"Attributed Question and Answering (AQA) refers to answering questions grounded to a given corpus and providing citation\")])\n",
        "grounding_passages.passages.append(glm.GroundingPassage(content=passage_a, id=\"001\"))\n",
        "passage_b = glm.Content(parts=[glm.Part(text=\"An LLM is not designed to generate content grounded in a set of passages. Although instructing an LLM to answer questions only based on a set of passages reduces hallucination, hallucination still often occurs when LLMs generate responses unsupported by facts provided by passages\")])\n",
        "grounding_passages.passages.append(glm.GroundingPassage(content=passage_b, id=\"002\"))\n",
        "passage_c = glm.Content(parts=[glm.Part(text=\"Hallucination is one of the biggest problems in Large Language Models (LLM) development. Large Language Models (LLMs) could produce responses that are fictitious and incorrect, which significantly impacts the usefulness and trustworthiness of applications built with language models.\")])\n",
        "grounding_passages.passages.append(glm.GroundingPassage(content=passage_c, id=\"003\"))\n",
        "\n",
        "# Create the request\n",
        "req = glm.GenerateAnswerRequest(model=MODEL_NAME,\n",
        "                                contents=[user_query_content],\n",
        "                                inline_passages=grounding_passages,\n",
        "                                answer_style=answer_style)\n",
        "aqa_response = generative_service_client.generate_answer(req)\n",
        "print(aqa_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QMrv_LTJGKGz"
      },
      "source": [
        "## Share the corpus\n",
        "\n",
        "You can choose to share the corpus with others using the [`CreatePermissionRequest`](https://ai.google.dev/api/python/google/generativeai/protos/CreatePermissionRequest) API.\n",
        "\n",
        "Constraints:\n",
        "\n",
        "* There are 2 roles for sharing: `READER` and `EDITOR`.\n",
        "  * A `READER` can query the corpus.\n",
        "  * A `WRITER` has reader's permissions and additionally can edit and share the corpus.\n",
        "* A corpus can be public by granting `EVERYONE` as `user_type` read access."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VCV4UgqUQOdv"
      },
      "outputs": [],
      "source": [
        "# Replace your-email@gmail.com with the email added as a test user in the OAuth Quickstart\n",
        "shared_user_email = \"TODO-your-email@gmail.com\" #  @param {type:\"string\"}\n",
        "user_type = \"USER\"\n",
        "role = \"READER\"\n",
        "\n",
        "# Make the request\n",
        "# corpus_resource_name is a variable set in the \"Create a corpus\" section.\n",
        "request = glm.CreatePermissionRequest(\n",
        "    parent=corpus_resource_name,\n",
        "    permission=glm.Permission(grantee_type=user_type,\n",
        "                              email_address=shared_user_email,\n",
        "                              role=role))\n",
        "create_permission_response = permission_service_client.create_permission(request)\n",
        "print(create_permission_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n7jUXKkZVOjn"
      },
      "source": [
        "## Delete the corpus\n",
        "\n",
        "Use [`DeleteCorpusRequest`](https://ai.google.dev/api/python/google/generativeai/protos/DeleteCorpusRequest) to delete a user corpus and all associated `Document`s & `Chunk`s.\n",
        "\n",
        "Note that non-empty corpora will throw an error without specifying an `force=True` flag. If you set `force=True`, any `Chunk`s and objects related to this `Document` will also be deleted.\n",
        "\n",
        "If `force=False` (the default) and the `Document` contains any `Chunk`s, a `FAILED_PRECONDITION` error will be returned."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yAKGJwrM0Zs8"
      },
      "outputs": [],
      "source": [
        "# Set force to False if you don't want to delete non-empty corpora.\n",
        "req = glm.DeleteCorpusRequest(name=corpus_resource_name, force=True)\n",
        "delete_corpus_response = retriever_service_client.delete_corpus(req)\n",
        "print(\"Successfully deleted corpus: \" + corpus_resource_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b8kIyLsUVC_Y"
      },
      "source": [
        "## Summary and further reading\n",
        "\n",
        "This guide introduced the Semantic Retriever and Attributed Question & Answering (AQA) APIs of the Gemini API and showed how you can use it to perform semantic information retrieval on your custom text data. Note that this API also works with the [LlamaIndex](https://www.llamaindex.ai/){:.external} data framework. Refer to [the tutorial](https://github.com/run-llama/llama_index/blob/main/docs/docs/examples/managed/GoogleDemo.ipynb){:.external} to learn more.\n",
        "\n",
        "Also refer to the [API docs](https://ai.google.dev/api) to learn more about the other available functionalities.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9YGv4x9ehLba"
      },
      "source": [
        "## Appendix: Setup OAuth with user credentials {:#user-oauth}\n",
        "\n",
        "Follow the steps below from the [OAuth Quickstart](https://ai.google.dev/docs/oauth_quickstart) to setup OAuth authentication.\n",
        "\n",
        "1. [Configure the OAuth consent screen](https://ai.google.dev/docs/oauth_quickstart#configure-oauth).\n",
        "\n",
        "1. [Authorize credentials for a desktop application](https://ai.google.dev/docs/oauth_quickstart#authorize-credentials). To run this notebook in Colab, first rename your credential file (usually `client_secret_*.json`) to just `client_secret.json`. Then upload the file by using the file icon on the left sidebar, then the upload icon, as shown in the screenshot below.\n",
        "\n",
        "<img width=400 src=\"https://ai.google.dev/tutorials/images/colab_upload.png\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9C3X6r1dueO4"
      },
      "outputs": [],
      "source": [
        "# Replace TODO-your-project-name with the project used in the OAuth Quickstart\n",
        "project_name = \"TODO-your-project-name\" #  @param {type:\"string\"}\n",
        "# Replace TODO-your-email@gmail.com with the email added as a test user in the OAuth Quickstart\n",
        "email = \"TODO-your-email@gmail.com\" #  @param {type:\"string\"}\n",
        "# Rename the uploaded file to `client_secret.json` OR\n",
        "# Change the variable `client_file_name` in the code below.\n",
        "client_file_name = \"client_secret.json\"\n",
        "\n",
        "# IMPORTANT: Follow the instructions from the output - you must copy the command\n",
        "# to your terminal and copy the output after authentication back here.\n",
        "!gcloud config set project $project_name\n",
        "!gcloud config set account $email\n",
        "\n",
        "# NOTE: The simplified project setup in this tutorial triggers a \"Google hasn't verified this app.\" dialog.\n",
        "# This is normal, click \"Advanced\" -> \"Go to [app name] (unsafe)\"\n",
        "!gcloud auth application-default login --no-browser --client-id-file=$client_file_name --scopes=\"https://www.googleapis.com/auth/generative-language.retriever,https://www.googleapis.com/auth/cloud-platform\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lybTOlN9gz0I"
      },
      "source": [
        "Initialize the client library and re-run the notebook starting from [Create a corpus](#create-corpus)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9MJbDTpMmlKJ"
      },
      "outputs": [],
      "source": [
        "import google.ai.generativelanguage as glm\n",
        "\n",
        "generative_service_client = glm.GenerativeServiceClient()\n",
        "retriever_service_client = glm.RetrieverServiceClient()\n",
        "permission_service_client = glm.PermissionServiceClient()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "semantic_retrieval.ipynb",
      "toc_visible": true
    },
    "google": {
      "image_path": "/site-assets/images/share.png",
      "keywords": [
        "examples",
        "googleai",
        "samplecode",
        "python",
        "embed",
        "aqa"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
